import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

# ------------------------------
# 传感器坐标生成 (保持不变)
# ------------------------------
def generate_sensor_positions():
    angles = np.deg2rad([0, 45, 90, 135, 180, 225, 270, 315])
    radii = [45000, 60000, 72500, 120000]  # 内圈到外圈的半径比例
    positions = {}
    
    for angle in angles:
        deg = int(np.rad2deg(angle))
        for j, r in enumerate(radii):
            x = r * np.cos(angle)
            y = r * np.sin(angle)
            positions[f"In{deg}_{j+1}"] = (x, y)
            if j == 3:  # 最外圈的传感器标记为Mid
                positions[f"Mid{deg}"] = (x, y)
    return positions

# ------------------------------
# 平面拟合函数（保持不变）
# ------------------------------
def fit_plane1(data, sensor_pos):
    """使用最小二乘法拟合平面 z = a*x + b*y + c"""
    points = []
    for name, value in data.items():
        if name in sensor_pos:
            x, y = sensor_pos[name]
            points.append([x, y, value])
    
    if not points:
        raise ValueError("没有找到匹配的传感器数据")
    
    points = np.array(points)
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]
    
    A = np.column_stack([x, y, np.ones_like(x)])
    
    coeffs, _, _, _ = np.linalg.lstsq(A, z, rcond=None)
    a, b, c = coeffs
    
    residuals = z - (a*x + b*y + c)
    rms_error = np.sqrt(np.mean(residuals**2))
    
    return a, b, c, rms_error


def fit_plane(data, sensor_pos):
    """增强健壮性的平面拟合函数"""
    points = []
    
    # 1. 验证输入类型
    if not isinstance(data, dict) or not isinstance(sensor_pos, dict):
        raise TypeError("输入必须是字典类型")
    
    # 2. 安全处理数据
    for name, value in data.items():
        try:
            if name is None:
                continue
                
            pos = sensor_pos.get(name)
            if pos is None:
                continue
                
            x, y = pos
            points.append([x, y, value])
        except (TypeError, ValueError) as e:
            print(f"跳过无效数据点 {name}: {e}")
            continue
    
    if len(points) < 3:
        raise ValueError("至少需要3个有效点进行平面拟合")
    
    # 3. 转换为numpy数组
    import numpy as np
    points = np.array(points)
    x = points[:, 0]
    y = points[:, 1]
    z = points[:, 2]
    
    # 4. 平面拟合
    A = np.column_stack([x, y, np.ones_like(x)])
    coeffs, _, _, _ = np.linalg.lstsq(A, z, rcond=None)
    
    # 5. 计算误差
    residuals = z - (coeffs[0]*x + coeffs[1]*y + coeffs[2])
    rms_error = np.sqrt(np.mean(residuals**2))
    
    return coeffs[0], coeffs[1], coeffs[2], rms_error
# ------------------------------
# 可视化函数（保持不变）
# ------------------------------
def plot_plane_and_data(ax, data, sensor_pos, a, b, c, title, color='blue'):
    """绘制传感器数据和拟合平面"""
    points = []
    for name, value in data.items():
        if name in sensor_pos:
            x, y = sensor_pos[name]
            points.append([x, y, value])
    
    points = np.array(points)
    x_data = points[:, 0]
    y_data = points[:, 1]
    z_data = points[:, 2]
    
    ax.scatter(x_data, y_data, z_data, c=color, s=50, depthshade=True, 
                        label=f'{title} Data')
    
    x_min, x_max = min(x_data)-1, max(x_data)+1
    y_min, y_max = min(y_data)-1, max(y_data)+1
    xx, yy = np.meshgrid(np.linspace(x_min, x_max, 20),
                         np.linspace(y_min, y_max, 20))
    zz = a * xx + b * yy + c
    
    ax.plot_surface(xx, yy, zz, alpha=0.5, color=color)
    
    from matplotlib.patches import Patch
    legend_elements = [
        Patch(facecolor=color, alpha=0.5, label=f'{title} Plane'),
        Patch(facecolor=color, label=f'{title} Data')
    ]
    ax.legend(handles=legend_elements)
    
    ax.set_xlabel('X Position')
    ax.set_ylabel('Y Position')
    ax.set_zlabel('Z Value')
    ax.set_title(f'{title} Plane Fitting')

# ------------------------------
# 计算调平参数（保持不变）
# ------------------------------
def calculate_adjustment(a_top, b_top, a_bottom, b_bottom):
    """计算需要调整的rx和ry（微弧度）"""
    rx_rad = np.arctan(a_bottom) - np.arctan(a_top)
    ry_rad = np.arctan(b_bottom) - np.arctan(b_top)
    
    rx_µrad = rx_rad * 1e6
    ry_µrad = ry_rad * 1e6
    
    return rx_µrad, ry_µrad

# ------------------------------
# 核心计算函数，供C#调用
# ------------------------------
def process_sensor_data(top_data, bottom_data, adjusted_data, save_path=None):
    """
    处理传感器数据并显示图表窗口
    参数:
        save_path: 图片保存路径（可选）
    """
    # 生成传感器位置（与之前相同）
    sensor_pos = generate_sensor_positions()  
    
    try:
        # 平面拟合计算
        a_top, b_top, c_top, _ = fit_plane(top_data, sensor_pos)
        a_bottom, b_bottom, c_bottom, _ = fit_plane(bottom_data, sensor_pos)
        a_adj, b_adj, c_adj, _ = fit_plane(adjusted_data, sensor_pos)

        # 创建图表窗口
        fig = plt.figure(figsize=(18, 6), num="Sensor Data Visualization")
        
        # 子图1: Topchuck
        ax1 = fig.add_subplot(131, projection='3d')
        plot_plane_and_data(ax1, top_data, sensor_pos, a_top, b_top, c_top, 'Topchuck', 'blue')
        
        # 子图2: Bottomchuck
        ax2 = fig.add_subplot(132, projection='3d')
        plot_plane_and_data(ax2, bottom_data, sensor_pos, a_bottom, b_bottom, c_bottom, 'Bottomchuck', 'red')
        
        # 子图3: Adjusted
        ax3 = fig.add_subplot(133, projection='3d')
        plot_plane_and_data(ax3, adjusted_data, sensor_pos, a_adj, b_adj, c_adj, 'Adjusted', 'green')

        plt.tight_layout()
        
        # 保存图片（如果指定路径）
        if save_path:
            dirname = os.path.dirname(save_path)
            if dirname: os.makedirs(dirname, exist_ok=True)
            plt.savefig(save_path, dpi=150, bbox_inches='tight')
            print(f"图表已保存到: {save_path}")
        
        # 显示窗口（阻塞模式）
        # plt.show(block=True)  # 关键点：block=True保持窗口打开

        # input("按Enter键继续...")  # 保持窗口打开
        plt.close()
        # 返回计算结果
        return {
            "top_plane": {"a": a_top, "b": b_top, "c": c_top},
            "bottom_plane": {"a": a_bottom, "b": b_bottom, "c": c_bottom},
            "adjusted_plane": {"a": a_adj, "b": b_adj, "c": c_adj}
        }
        
    except Exception as e:
        return {"error": str(e)}
    finally:
        plt.close('all')  # 确保释放资源
# 注意：这里移除了 if __name__ == "__main__": 部分，因为Python.NET会直接调用 process_sensor_data 函数。
# 如果你还需要独立运行这个脚本进行测试，可以添加回一个简单的调用。
